import os
import cv2
import numpy as np

# 路径配置
images_dir = "E:/Python/PythonProject/detect/results"  # 图片目录
labels_dir = "E:/Python/PythonProject/detect/labels/train"  # 标签目录
output_dir = "E:/Python/PythonProject/detect/visualized"  # 可视化结果保存目录（可选）
os.makedirs(output_dir, exist_ok=True)  # 创建输出目录

# 类别列表（根据你的数据集修改）
class_names = ["class0", "class1", "class2","class3"]  # 替换为你的实际类别

# 颜色列表（每个类别一个颜色）
colors = np.random.randint(0, 255, size=(len(class_names), 4), dtype=np.uint8)

def plot_boxes(image, boxes, class_ids):
    """在图片上绘制边界框和标签"""
    h, w = image.shape[:2]
    for box, class_id in zip(boxes, class_ids):
        # 解析YOLO格式（归一化坐标 -> 绝对坐标）
        x_center, y_center, width, height = box
        x1 = int((x_center - width / 2) * w)
        y1 = int((y_center - height / 2) * h)
        x2 = int((x_center + width / 2) * w)
        y2 = int((y_center + height / 2) * h)

        # 绘制矩形和标签
        color = tuple(map(int, colors[class_id]))
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        label = f"{class_names[class_id]}"
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

def visualize_labels(image_path, label_path, save_path=None):
    """可视化单个图片的标签"""
    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: 无法读取图片 {image_path}")
        return

    # 读取标签
    boxes = []
    class_ids = []
    if os.path.exists(label_path):
        with open(label_path, 'r') as f:
            for line in f.readlines():
                parts = line.strip().split()
                if len(parts) == 5:
                    class_id = int(parts[0])
                    box = list(map(float, parts[1:5]))
                    boxes.append(box)
                    class_ids.append(class_id)

    # 绘制边界框
    plot_boxes(image, boxes, class_ids)

    # 显示或保存结果
    if save_path:
        cv2.imwrite(save_path, image)
    else:
        cv2.imshow("Visualization", image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

# 遍历所有标签文件并可视化
for images_file in os.listdir(images_dir):
    if images_file.endswith(".jpg"):
        # 构建对应的图片路径（支持.jpg/.png等格式）
        label_name = os.path.splitext(images_file)[0]
        label_path = os.path.join(labels_dir, label_name + ".txt")
        image_path = os.path.join(images_dir, images_file)

        save_path = os.path.join(output_dir, f"vis_{label_name}.jpg")  # 保存可视化结果
        visualize_labels(image_path, label_path, save_path)